package com.hc.app.dto;

import com.hc.app.exception.FrequencyControlErrorEnum;
import com.hc.app.exception.FrequencyControlException;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.locks.ReentrantLock;

/**
 * @Author: huangcheng
 * @CreateTime: 2024-07-25  12:46
 * @Description: 令牌桶算法
 * @Version: 1.0
 */
@Slf4j
@Data
public class TokenBucketDTO extends FrequencyControlDto{
    /**
     * 令牌桶容量
     */
    private final long capacity;
    /**
     * 每秒补充的令牌数
     */
    private final double refillRate;
    /**
     * 当前令牌数量
     */
    private double tokens;
    /**
     * 上次补充令牌的时间
     */
    private long lastRefillTime;

    private final ReentrantLock lock = new ReentrantLock();

    /**
     * 构造初始容量
     */
    public TokenBucketDTO(long capacity, double refillRate) {
        if (capacity <= 0 || refillRate <= 0) {
            throw new FrequencyControlException(FrequencyControlErrorEnum.CAPACITY_REFILL_ERROR);
        }
        this.capacity = capacity;
        this.refillRate = refillRate;
        this.tokens = capacity;
        this.lastRefillTime = System.nanoTime();
    }

    /**
     * 是否获取令牌
     * @param permits
     * @return
     */
    public boolean tryAcquire(int permits) {
        lock.lock();
        try {
            refillTokens();
            if (tokens < permits) {
                return true;
            }
            return false;
        } finally {
            lock.unlock();
        }
    }

    /**
     * 使用令牌
     * @param permits
     */
    public void deductionToken(int permits) {
        lock.lock();
        try {
            tokens -= permits;
        } finally {
            lock.unlock();
        }
    }

    /**
     * 补充令牌：count = 这次和上次时间差*每秒补充个数
     */
    private void refillTokens() {
        long currentTime = System.nanoTime();
        // 转换为秒
        double elapsedTime = (currentTime - lastRefillTime) / 1e9;
        double tokensToAdd = elapsedTime * refillRate;
        log.info("tokensToAdd is {}", tokensToAdd);
        // 令牌总数不能超过令牌桶容量
        tokens = Math.min(capacity, tokens + tokensToAdd);
        log.info("current tokens is {}", tokens);
        lastRefillTime = currentTime;
    }
}
